import math
import rclpy
from rclpy.lifecycle import LifecycleNode
from rclpy.action import ActionClient
from geometry_msgs.msg import PoseStamped, Quaternion
from std_msgs.msg import Float64MultiArray, Bool
from nav_msgs.msg import Path, Odometry
from rclpy.lifecycle import Publisher
import numpy as np
from rclpy.qos import DurabilityPolicy, QoSProfile
from message_filters import Subscriber, TimeSynchronizer, ApproximateTimeSynchronizer
from general_interface.action import SwitchConnection, SwitchMode
from general_interface.msg import WorkMode
import tf_transformations as tft


class AutoDockTask(LifecycleNode):

    def __init__(self):
        super().__init__("auto_dock")
        self.declare_parameter("robot1", "robot1")
        self.declare_parameter("robot2", "robot2")
        self.declare_parameter("target1", "odom")
        self.declare_parameter("target2", "odom")
        self.declare_parameter("radius", 0.2)
        self.declare_parameter("step_length", 0.1)
        self.declare_parameter("start_vec", 1.0)
        self.declare_parameter("dock_vec", 0.1)

    def on_configure(self, state):
        self.robot1: str = self.get_parameter("robot1").value
        self.robot2: str = self.get_parameter("robot2").value
        self.target1: str = self.get_parameter("target1").value
        self.target2: str = self.get_parameter("target2").value
        self.radius: float = self.get_parameter("radius").value
        self.stepLength: float = self.get_parameter("step_length").value
        self.startVec: float = self.get_parameter("start_vec").value
        self.dockVec: float = self.get_parameter("dock_vec").value

        qos = QoSProfile(depth=1, durability=DurabilityPolicy.TRANSIENT_LOCAL)
        self.path_pub1: Publisher = self.create_publisher(
            Path, f"/{self.robot1}/reference_path", qos
        )
        self.vec_pub1: Publisher = self.create_publisher(
            Float64MultiArray, f"/{self.robot1}/reference_velocity", qos
        )
        self.coreClient1 = ActionClient(self, SwitchMode, f"/{self.robot1}/switch_mode")
        self.path_pub2: Publisher = self.create_publisher(
            Path, f"/{self.robot2}/reference_path", qos
        )
        self.vec_pub2: Publisher = self.create_publisher(
            Float64MultiArray, f"/{self.robot2}/reference_velocity", qos
        )
        self.coreClient2 = ActionClient(self, SwitchMode, f"/{self.robot2}/switch_mode")
        self.connectorClient = ActionClient(
            self, SwitchConnection, f"/{self.robot1}/switch_connection"
        )
        self.first_time = True
        return super().on_configure(state)

    def on_cleanup(self, state):
        self.destroy_publisher(self.path_pub1)
        self.destroy_publisher(self.vec_pub1)
        self.destroy_publisher(self.path_pub2)
        self.destroy_publisher(self.vec_pub2)
        self.destroy_client(self.connectorClient)
        return super().on_cleanup(state)

    def on_activate(self, state):
        self.get_logger().info("waiting for server")
        self.coreClient1.wait_for_server(1)
        self.coreClient2.wait_for_server(1)
        self.get_logger().info("wait server finished")
        goal = SwitchMode.Goal()
        goal.mode.mode = WorkMode.PATHFOLLOW
        future1 = self.coreClient1.send_goal_async(goal)
        future2 = self.coreClient2.send_goal_async(goal)
        # rclpy.spin_until_future_complete(self, future1)
        # rclpy.spin_until_future_complete(self, future2)
        # self.get_logger().info("wait for state change for core")

        odomSub1 = Subscriber(self, Odometry, f"/{self.target1}")
        odomSub2 = Subscriber(self, Odometry, f"/{self.target2}")
        self.odomTs = ApproximateTimeSynchronizer([odomSub1, odomSub2], 100, slop=0.1)
        self.odomTs.registerCallback(self.odom_callback)
        goal = SwitchConnection.Goal()
        goal.command = SwitchConnection.Goal.CONNECT
        self.future_done = self.connectorClient.send_goal_async(goal)
        self.future_done.add_done_callback(self.goal_response_callback)
        self.get_logger().info("auto docking start, task node activated")
        return super().on_activate(state)

    def on_deactivate(self, state):
        self.odomTs = None
        self.get_logger().info("auto docking task node deactivated")
        return super().on_deactivate(state)

    def odom_callback(self, start: Odometry, end: Odometry):
        #self.get_logger().info("odom callback")
        path1, vec1 = self.solution1(
            start,
            end,
            self.radius,
            self.startVec,
            self.dockVec,
        )
        self.path_pub1.publish(path1)
        self.vec_pub1.publish(vec1)

        path2, vec2 = self.solution1(
            end,
            start,
            self.radius,
            -self.startVec,
            -self.dockVec,
        )
        self.path_pub2.publish(path2)
        self.vec_pub2.publish(vec2)

    def goal_response_callback(self, future):
        goal_handle = future.result()
        if not goal_handle.accepted:
            self.get_logger().info("Goal rejected :(")
            return
        self.future_result = goal_handle.get_result_async()
        self.future_result.add_done_callback(self.result_callback)

    def result_callback(self, future):
        result = future.result().result
        self.get_logger().info("Result: {0}".format(result))
        self.get_logger().info("Goal finished, please deactivate")

        goal = SwitchMode.Goal()
        goal.mode.mode = WorkMode.IDLE
        future1 = self.coreClient1.send_goal_async(goal)
        future2 = self.coreClient2.send_goal_async(goal)
        rclpy.spin_until_future_complete(self, future1)
        rclpy.spin_until_future_complete(self, future2)

    def connection_callback(self, msg: Bool):
        if msg.data:
            self.get_logger().info("connection finished")

    def solution1(
        self,
        start: Odometry,
        end: Odometry,
        l: float,
        start_vel: float,
        dock_vec: float,
    ) -> tuple[Path, Float64MultiArray]:
        """
        起点和终点以及初始朝向的交点构成三角形，实际轨迹为与三角形相似，但形状更小
        """
        startP = np.array([start.pose.pose.position.x, start.pose.pose.position.y])
        endP = np.array([end.pose.pose.position.x, end.pose.pose.position.y])

        distance = np.linalg.norm(endP - startP)
        l = distance * l

        startQua = (
            start.pose.pose.orientation.x,
            start.pose.pose.orientation.y,
            start.pose.pose.orientation.z,
            start.pose.pose.orientation.w,
        )
        startDir = tft.quaternion_matrix(startQua)[0:2, 0]
        startDir /= np.linalg.norm(startDir)
        endQua = (
            end.pose.pose.orientation.x,
            end.pose.pose.orientation.y,
            end.pose.pose.orientation.z,
            end.pose.pose.orientation.w,
        )
        endDir = tft.quaternion_matrix(endQua)[0:2, 0]
        endDir /= np.linalg.norm(endDir)

        if start_vel < 0:
            startMid = startP - startDir * l
            endMid = endP + endDir * l
            midDir = endMid - startMid
            midDir /= np.linalg.norm(midDir)
            startMidInLine = startMid + midDir * l
            endMidInLine = endMid - midDir * l
            mid = (startMidInLine + endMidInLine) / 2
        else:
            startMid = startP + startDir * l
            endMid = endP - endDir * l
            midDir = endMid - startMid
            midDir /= np.linalg.norm(midDir)
            startMidInLine = startMid + midDir * l
            endMidInLine = endMid - midDir * l
            mid = (startMidInLine + endMidInLine) / 2

        startPoints = self.besizer_curve(
            startP,
            startMid,
            startMidInLine,
            int(np.linalg.norm(startP - startMidInLine) / self.stepLength),
        )
        endPoints = self.besizer_curve(
            endMidInLine,
            endMid,
            endP,
            int(np.linalg.norm(endMidInLine - endP) / self.stepLength),
        )
        midPoints = self.besizer_curve(
            startMidInLine,
            mid,
            endMidInLine,
            int(np.linalg.norm(startMidInLine - endMidInLine) / self.stepLength),
        )
        points = np.vstack([startPoints, midPoints[1:-1, :], endPoints])
        velocity = np.linspace(start_vel, dock_vec, len(points))

        path = Path()
        path.header.frame_id = "map"
        path.header.stamp = self.get_clock().now().to_msg()
        for i in range(len(points)):
            pose = PoseStamped()
            pose.header.frame_id = "map"
            pose.header.stamp = self.get_clock().now().to_msg()
            pose.pose.position.x = points[i][0]
            pose.pose.position.y = points[i][1]
            pose.pose.position.z = 0.0
            if i == len(points) - 1:
                pose.pose.orientation = end.pose.pose.orientation
            else:
                dx = points[i + 1][0] - points[i][0]
                dy = points[i + 1][1] - points[i][1]
                # 插值计算点的朝向
                qua = tft.quaternion_from_euler(
                    0.0,
                    0.0,
                    (
                        math.atan2(
                            -dy,
                            -dx,
                        )
                        if start_vel < 0
                        else math.atan2(dy, dx)
                    ),
                )
                pose.pose.orientation = Quaternion(
                    x=qua[0], y=qua[1], z=qua[2], w=qua[3]
                )
            path.poses.append(pose)
        vec = Float64MultiArray()
        vec.data = velocity
        return path, vec

    def besizer_curve(self, p0, p1, p2, num_points):
        t = np.linspace(0, 1, num_points)[:, np.newaxis]
        points = (1 - t) ** 2 * p0 + 2 * t * (1 - t) * p1 + t**2 * p2
        return points


def main(args=None):
    rclpy.init(args=args)
    task = AutoDockTask()
    task.trigger_configure()
    rclpy.spin(task)
    rclpy.shutdown()


if __name__ == "__main__":
    main()
